{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "from matplotlib import pyplot as plt\n",
    "import string\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams[\"font.size\"] = 16\n",
    "SHOW_RG = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"pg_run.txt\") as f:\n",
    "    data = f.read().split(\"\\n\")[2:]\n",
    "data = [x.split(\" \") for x in data if len(x) > 1 and (x[0] in string.digits or x[0] == \"x\")]\n",
    "\n",
    "data = [(x[0], x[1], float(x[2]), x[3], float(x[4])) for x in data]\n",
    "pg_data = data\n",
    "pg_times = np.array([x[2] for x in pg_data])\n",
    "pg_times -= np.min(pg_times)\n",
    "pg_times /= 60\n",
    "\n",
    "\n",
    "def read_bao_data(fp):\n",
    "    with open(fp) as f:\n",
    "        data = f.read().split(\"\\n\")[2:]\n",
    "\n",
    "    training_times = []\n",
    "    for idx in range(len(data)):\n",
    "        if data[idx].strip().startswith(\"Initial input channels\"):\n",
    "            prev_line = data[idx-1].split(\" \")\n",
    "            if prev_line[0] == \"Retry\":\n",
    "                continue\n",
    "            training_times.append(float(prev_line[2]))\n",
    "\n",
    "\n",
    "    training_times = np.array(training_times)\n",
    "\n",
    "    data = [x.split(\" \") for x in data if len(x) > 1 and (x[0] in string.digits or x[0] == \"x\")]\n",
    "    data = [(x[0], x[1], float(x[2]), x[3], float(x[4])) for x in data]\n",
    "    bao_data = data\n",
    "\n",
    "    bao_times = np.array([x[2] for x in bao_data])\n",
    "    training_times -= np.min(bao_times)\n",
    "    bao_times -= np.min(bao_times)\n",
    "\n",
    "    bao_times /= 60\n",
    "    training_times /= 60\n",
    "    return bao_data, bao_times, training_times\n",
    "\n",
    "bao_data, bao_times, training_times = read_bao_data(\"bao_run.txt\")\n",
    "if SHOW_RG:\n",
    "    bao_rb_data, bao_rb_times, training_rb_times = read_bao_data(\"bao_with_regblock.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "queries_complete = np.arange(0, len(pg_times))\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, constrained_layout=True)\n",
    "\n",
    "\n",
    "train_y = []\n",
    "train_rb_y = []\n",
    "for tt in training_times:\n",
    "    idx = np.searchsorted(bao_times, tt)\n",
    "    train_y.append(idx)\n",
    "    \n",
    "if SHOW_RG:\n",
    "    for tt in training_rb_times:\n",
    "        idx = np.searchsorted(bao_rb_times, tt)\n",
    "        train_rb_y.append(idx)\n",
    "    \n",
    "plt.scatter(training_times, train_y, s=45, color=\"red\", label=\"Training\")\n",
    "\n",
    "ax.plot(pg_times, queries_complete, label=\"PostgreSQL\", lw=3)\n",
    "ax.plot(bao_times, queries_complete, label=\"Bao\", lw=3)\n",
    "\n",
    "if SHOW_RG:\n",
    "    plt.scatter(training_rb_times, train_rb_y, s=45, color=\"red\")\n",
    "    ax.plot(bao_rb_times, queries_complete, label=\"Bao (w/ exploration)\", lw=3)\n",
    "\n",
    "ax.set_xlabel(\"Time (m)\")\n",
    "ax.set_ylabel(\"Queries complete\")\n",
    "ax.set_title(\"PostgreSQL vs Bao Optimizer\")\n",
    "\n",
    "ax.grid(linestyle=\"--\", linewidth=1)\n",
    "ax.legend()\n",
    "fig.savefig(\"queries_vs_time.svg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_pg_times = sorted([x[4] for x in pg_data])\n",
    "all_bao_times = sorted([x[4] for x in bao_data])\n",
    "\n",
    "if SHOW_RG:\n",
    "    all_bao_rb_times = sorted([x[4] for x in bao_rb_data])\n",
    "\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4), constrained_layout=True)\n",
    "\n",
    "ax = axes[0]\n",
    "ax.plot(np.linspace(0, 1, len(all_pg_times)), all_pg_times, lw=3, label=\"PostgreSQL\")\n",
    "ax.plot(np.linspace(0, 1, len(all_pg_times)), all_bao_times, lw=3, label=\"Bao\")\n",
    "\n",
    "if SHOW_RG:\n",
    "    ax.plot(np.linspace(0, 1, len(all_pg_times)), all_bao_rb_times, lw=3, label=\"Bao (w/ exploration)\")\n",
    "\n",
    "ax.grid(linestyle=\"--\", linewidth=1)\n",
    "ax.set_xlabel(\"Proportion of Queries\")\n",
    "ax.set_ylabel(\"Max Latency (s)\")\n",
    "ax.set_title(\"Query Latency CDF\")\n",
    "ax.legend()\n",
    "#ax.set_yscale(\"log\")\n",
    "\n",
    "\n",
    "ax = axes[1]\n",
    "ax.plot(np.linspace(0, 1, len(all_pg_times)), all_pg_times, lw=3, label=\"PostgreSQL\")\n",
    "ax.plot(np.linspace(0, 1, len(all_pg_times)), all_bao_times, lw=3, label=\"Bao\")\n",
    "\n",
    "if SHOW_RG:\n",
    "    ax.plot(np.linspace(0, 1, len(all_pg_times)), all_bao_rb_times, lw=3, label=\"Bao (w/ exploration)\")\n",
    "\n",
    "ax.grid(linestyle=\"--\", linewidth=1)\n",
    "ax.set_xlabel(\"Proportion of Queries\")\n",
    "ax.set_ylabel(\"Max Latency (s)\")\n",
    "ax.set_title(\"Query Latency CDF (log scale)\")\n",
    "ax.legend()\n",
    "ax.set_yscale(\"log\")\n",
    "fig.savefig(\"cdf.svg\")\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the last PG time for each query\n",
    "pg_query_time = {}\n",
    "for itm in pg_data:\n",
    "    pg_query_time[itm[3]] = itm[4]\n",
    "        \n",
    "# get each Bao time\n",
    "bao_query_times = defaultdict(list)\n",
    "for itm in bao_data[50:]:\n",
    "    bao_query_times[itm[3]].append(itm[4])\n",
    "    \n",
    "if SHOW_RG:\n",
    "    # get each Bao time\n",
    "    bao_rb_query_times = defaultdict(list)\n",
    "    for itm in bao_rb_data[50:]:\n",
    "        bao_rb_query_times[itm[3]].append(itm[4])\n",
    "\n",
    "max_repeats = max(len(x) for x in bao_query_times.values())\n",
    "\n",
    "def extract_q_number(x):\n",
    "    return int(x[x.find(\"/q\")+2:x.find(\"_\", x.find(\"/q\"))])\n",
    "\n",
    "q_order = sorted(bao_query_times.keys(), key=extract_q_number)\n",
    "\n",
    "grid = [bao_query_times[x] for x in q_order]\n",
    "\n",
    "if SHOW_RG:\n",
    "    grid_rb = [bao_rb_query_times[x] for x in q_order]\n",
    "\n",
    "\n",
    "reg_data = []\n",
    "for idx, q in enumerate(q_order):\n",
    "    if SHOW_RG:\n",
    "        reg_data.append({\"Q\": f\"q{extract_q_number(q)}\", \n",
    "                         \"PG\": pg_query_time[q],\n",
    "                         \"Bao worst\": max(grid[idx]),\n",
    "                         \"Bao best\": min(grid[idx]),\n",
    "                         \"Bao + E worst\": max(grid_rb[idx]),\n",
    "                         \"Bao + E best\": min(grid_rb[idx])})\n",
    "    else:\n",
    "        reg_data.append({\"Q\": f\"q{extract_q_number(q)}\", \n",
    "                 \"PG\": pg_query_time[q],\n",
    "                 \"Bao worst\": max(grid[idx]),\n",
    "                 \"Bao best\": min(grid[idx])})\n",
    "                \n",
    "\n",
    "\n",
    "def color_regression(col):\n",
    "    def c_for_diff(diff):\n",
    "        if diff < 2 and diff > -2:\n",
    "            return \"background-color: white\"\n",
    "        elif diff > 0.5:\n",
    "            return \"background-color: #f27281\"\n",
    "        else:\n",
    "            return \"background-color: #9ee3ad\"\n",
    "        \n",
    "    to_r = [\"\"]\n",
    "    \n",
    "    if SHOW_RG:\n",
    "        pg, bao_worst, bao_best, bao_rg_worst, bao_rg_best = col\n",
    "    else:\n",
    "        pg, bao_worst, bao_best = col\n",
    "\n",
    "    \n",
    "    to_r.append(c_for_diff(bao_worst - pg))\n",
    "    to_r.append(c_for_diff(bao_best - pg))\n",
    "    \n",
    "    if SHOW_RG:\n",
    "        to_r.append(c_for_diff(bao_rg_worst - pg))    \n",
    "        to_r.append(c_for_diff(bao_rg_best - pg))\n",
    "\n",
    "    return to_r\n",
    "\n",
    "reg_data = pd.DataFrame(reg_data).set_index(\"Q\")\n",
    "reg_data.style.apply(color_regression, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
